// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/sparksql/array_sort.h>

#include <memory>
#include <pollux/common/base/exceptions.h>
#include <pollux/expression/eval_ctx.h>
#include <pollux/expression/expr.h>
#include <pollux/expression/vector_function.h>
#include <pollux/functions/lib/rows_translation_util.h>
#include <pollux/functions/sparksql/comparisons.h>
#include <pollux/type/type.h>
#include <pollux/vector/base_vector.h>
#include <pollux/vector/complex_vector.h>
#include <pollux/vector/selectivity_vector.h>
#include <pollux/vector/simple_vector.h>
#include <pollux/vector/type_aliases.h>

namespace kumo::pollux::functions::sparksql {
    namespace {
        void applyComplexType(
            const SelectivityVector &rows,
            ArrayVector *inputArray,
            bool ascending,
            bool nullsFirst,
            exec::EvalCtx &context,
            VectorPtr *resultElements) {
            auto elementsVector = inputArray->elements();

            // Allocate new vectors for indices.
            BufferPtr indices = allocate_indices(elementsVector->size(), context.pool());
            vector_size_t *rawIndices = indices->asMutable<vector_size_t>();

            const CompareFlags flags{.nullsFirst = nullsFirst, .ascending = ascending};
            // Note: Reusing offsets and sizes isn't safe if the input array had two
            // arrays that had overlapping (but not identical) ranges in the input.
            rows.applyToSelected([&](vector_size_t row) {
                auto size = inputArray->sizeAt(row);
                auto offset = inputArray->offsetAt(row);

                for (auto i = 0; i < size; ++i) {
                    rawIndices[offset + i] = offset + i;
                }
                std::sort(
                    rawIndices + offset,
                    rawIndices + offset + size,
                    [&](vector_size_t &a, vector_size_t &b) {
                        return elementsVector->compare(elementsVector.get(), a, b, flags) < 0;
                    });
            });

            *resultElements = BaseVector::transpose(indices, std::move(elementsVector));
        }

        template<typename T>
        inline void swapWithNull(
            FlatVector<T> *vector,
            vector_size_t index,
            vector_size_t nullIndex) {
            // Values are already present in vector stringBuffers. Don't create additional
            // copy.
            if constexpr (std::is_same_v<T, StringView>) {
                vector->setNoCopy(nullIndex, vector->value_at(index));
            } else {
                vector->set(nullIndex, vector->value_at(index));
            }
            vector->set_null(index, true);
        }

        template<TypeKind kind>
        void applyTyped(
            const SelectivityVector &rows,
            const ArrayVector *inputArray,
            bool ascending,
            bool nullsFirst,
            exec::EvalCtx &context,
            VectorPtr *resultElements) {
            using T = typename TypeTraits<kind>::NativeType;

            // Copy array elements to new vector.
            const VectorPtr &inputElements = inputArray->elements();
            SelectivityVector elementRows =
                    toElementRows(inputElements->size(), rows, inputArray);

            *resultElements = BaseVector::create(
                inputElements->type(), inputElements->size(), context.pool());
            (*resultElements)
                    ->copy(inputElements.get(), elementRows, /*toSourceRow=*/nullptr);

            auto flatResults = (*resultElements)->as_flat_vector<T>();
            T *resultRawValues = flatResults->mutableRawValues();

            auto processRow = [&](vector_size_t row) {
                auto size = inputArray->sizeAt(row);
                auto offset = inputArray->offsetAt(row);
                if (size == 0) {
                    return;
                }
                vector_size_t numNulls = 0;
                if (nullsFirst) {
                    // Move nulls to beginning of array.
                    for (vector_size_t i = 0; i < size; ++i) {
                        if (flatResults->is_null_at(offset + i)) {
                            swapWithNull<T>(flatResults, offset + numNulls, offset + i);
                            ++numNulls;
                        }
                    }
                } else {
                    // Move nulls to end of array.
                    for (vector_size_t i = size - 1; i >= 0; --i) {
                        if (flatResults->is_null_at(offset + i)) {
                            swapWithNull<T>(
                                flatResults, offset + size - numNulls - 1, offset + i);
                            ++numNulls;
                        }
                    }
                }
                // Exclude null values while sorting.
                auto rowBegin = offset + (nullsFirst ? numNulls : 0);
                auto rowEnd = rowBegin + size - numNulls;

                if constexpr (kind == TypeKind::BOOLEAN) {
                    uint64_t *rawBits = flatResults->template mutableRawValues<uint64_t>();
                    auto numSetBits = bits::countBits(rawBits, rowBegin, rowEnd);
                    // If ascending, false is placed before true, otherwise true is placed
                    // before false.
                    bool smallerValue = !ascending;
                    auto mid = ascending ? rowEnd - numSetBits : rowBegin + numSetBits;
                    bits::fillBits(rawBits, rowBegin, mid, smallerValue);
                    bits::fillBits(rawBits, mid, rowEnd, !smallerValue);
                } else {
                    if (ascending) {
                        std::sort(
                            resultRawValues + rowBegin, resultRawValues + rowEnd, Less<T>());
                    } else {
                        std::sort(
                            resultRawValues + rowBegin, resultRawValues + rowEnd, Greater<T>());
                    }
                }
            };

            rows.applyToSelected(processRow);
        }
    } // namespace

    void ArraySort::apply(
        const SelectivityVector &rows,
        std::vector<VectorPtr> &args,
        const TypePtr & /*outputType*/,
        exec::EvalCtx &context,
        VectorPtr &result) const {
        auto &arg = args[0];

        VectorPtr localResult;

        // Input can be constant or flat.
        if (arg->is_constant_encoding()) {
            auto *constantArray = arg->as<ConstantVector<ComplexType> >();
            const auto &flatArray = constantArray->value_vector();
            const auto flatIndex = constantArray->index();

            exec::LocalSingleRow singleRow(context, flatIndex);
            localResult = applyFlat(*singleRow, flatArray, context);
            localResult =
                    BaseVector::wrap_in_constant(rows.end(), flatIndex, localResult);
        } else {
            localResult = applyFlat(rows, arg, context);
        }

        context.moveOrCopyResult(localResult, rows, result);
    }

    VectorPtr ArraySort::applyFlat(
        const SelectivityVector &rows,
        const VectorPtr &arg,
        exec::EvalCtx &context) const {
        ArrayVector *inputArray = arg->as<ArrayVector>();
        VectorPtr resultElements;

        auto typeKind = inputArray->elements()->type_kind();
        if (typeKind == TypeKind::MAP || typeKind == TypeKind::ARRAY ||
            typeKind == TypeKind::ROW) {
            applyComplexType(
                rows, inputArray, ascending_, nullsFirst_, context, &resultElements);
        } else {
            POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH(
                applyTyped,
                typeKind,
                rows,
                inputArray,
                ascending_,
                nullsFirst_,
                context,
                &resultElements);
        }

        return std::make_shared<ArrayVector>(
            context.pool(),
            inputArray->type(),
            inputArray->nulls(),
            rows.end(),
            inputArray->offsets(),
            inputArray->sizes(),
            resultElements,
            inputArray->get_null_count());
    }

    // Signature: array_sort(array(T)) -> array(T)
    std::vector<std::shared_ptr<exec::FunctionSignature> > arraySortSignatures() {
        return {
            exec::FunctionSignatureBuilder()
            .typeVariable("T")
            .argumentType("array(T)")
            .returnType("array(T)")
            .build(),
        };
    }

    std::shared_ptr<exec::VectorFunction> makeArraySort(
        const std::string &name,
        const std::vector<exec::VectorFunctionArg> &inputArgs,
        const core::QueryConfig & /*config*/) {
        POLLUX_CHECK_EQ(inputArgs.size(), 1);
        // Nulls are considered largest.
        return std::make_shared<ArraySort>(/*ascending=*/true, /*nullsFirst=*/false);
    }

    // Signatures:
    //   sort_array(array(T)) -> array(T)
    //   sort_array(array(T), boolean) -> array(T)
    std::vector<std::shared_ptr<exec::FunctionSignature> > sortArraySignatures() {
        return {
            exec::FunctionSignatureBuilder()
            .typeVariable("T")
            .argumentType("array(T)")
            .returnType("array(T)")
            .build(),
            exec::FunctionSignatureBuilder()
            .typeVariable("T")
            .argumentType("array(T)")
            .constantArgumentType("boolean")
            .returnType("array(T)")
            .build(),
        };
    }

    std::shared_ptr<exec::VectorFunction> makeSortArray(
        const std::string &name,
        const std::vector<exec::VectorFunctionArg> &inputArgs,
        const core::QueryConfig & /*config*/) {
        POLLUX_CHECK(
            inputArgs.size() == 1 || inputArgs.size() == 2,
            "Invalid number of arguments {}, expected 1 or 2",
            inputArgs.size());
        bool ascending = true;
        // Read optional sort ascending flag.
        if (inputArgs.size() == 2) {
            BaseVector *boolVector = inputArgs[1].constantValue.get();
            if (!boolVector || !boolVector->is_constant_encoding()) {
                POLLUX_USER_FAIL(
                    "{} requires a constant bool as the second argument.", name);
            }
            ascending = boolVector->as<ConstantVector<bool> >()->value_at(0);
        }
        // Nulls are considered smallest.
        bool nullsFirst = ascending;
        return std::make_shared<ArraySort>(ascending, nullsFirst);
    }
} // namespace kumo::pollux::functions::sparksql
